# -*- coding: utf-8 -*-
import sys
import argparse
import arcpy
import numpy as np
from jenkspy import jenks_breaks

def jenkspy_Natural_Breaks_sample(inRaster_full_path, outRecalssRaster_full_path, nb_class, nb_times=50,
                                  sample_size=5000, nodata_val=0):
    """
    使用采样法对栅格进行自然断点分类，并正确处理NoData值。
    此版本通过强制转换NoData为NaN，解决了整型栅格NoData处理不当的问题。
    """
    arcpy.env.overwriteOutput = True
    try:
        # --- 第1步：获取输入栅格的NoData值（用于决定输出的NoData值）---
        inRas = arcpy.Raster(inRaster_full_path)
        desc = arcpy.Describe(inRas)
        input_nodata_val = desc.noDataValue
        arcpy.AddMessage("输入栅格的NoData值为: {}".format(input_nodata_val))

        # --- 第2步：将栅格转换为NumPy数组，并强制将所有NoData转为NaN ---
        # 这是解决问题的关键！无论输入栅格类型如何，这都会创建一个浮点型数组，
        # 其中所有NoData单元格都被明确地标记为NaN。
        arr = arcpy.RasterToNumPyArray(inRas, nodata_to_value=float('nan'))

        # 现在可以安全地使用 np.isnan 来创建掩码，分离有效数据和NoData
        mask = ~np.isnan(arr)
        arr_without_nan = arr[mask]  # 获取所有有效数据

        # --- 第3步：检查有效数据是否存在 ---
        if arr_without_nan.size == 0:
            arcpy.AddError("错误：输入栅格中没有有效数据，无法进行分类。")
            return None

        # --- 第4步：计算自然断点 ---
        group_breaks = []
        for i in range(nb_times):
            # 确保采样数不超过有效数据总数
            actual_sample_size = min(sample_size, arr_without_nan.size)
            tmp_arr = np.random.choice(arr_without_nan, actual_sample_size, replace=False)
            tmp_breaks = jenks_breaks(tmp_arr, nb_class=nb_class)
            group_breaks.append(tmp_breaks)

        group_breaks = np.array(group_breaks)
        break_values = group_breaks.mean(axis=0)

        # 微调断点以确保所有值都能被正确分类
        break_values[0] = np.min(arr_without_nan) - 0.001
        break_values[-1] = np.max(arr_without_nan) + 0.001

        arcpy.AddMessage("计算得到的自然断点值为：")
        arcpy.AddMessage(str(break_values))

        # --- 第5步：对数据进行分类并构建最终数组 ---
        classified_data = np.digitize(arr_without_nan, break_values, right=False)

        # 确定输出栅格的NoData值（优先继承输入的）
        output_nodata_val = input_nodata_val if input_nodata_val is not None else nodata_val

        # 创建一个与原始数组同样大小的浮点型最终输出数组，并用NaN填充
        final_arr = np.full(arr.shape, np.nan, dtype=np.float32)

        # 将分类好的数据填充到最终数组的有效数据位置
        final_arr[mask] = classified_data

        # 将原始NoData位置（现在是NaN）设置为我们确定的最终输出NoData值
        final_arr[~mask] = output_nodata_val

        # --- 第6步：将NumPy数组保存为栅格 ---
        output_cs = desc.spatialReference
        arcpy.env.outputCoordinateSystem = output_cs

        lowerLeft = arcpy.Point(inRas.extent.XMin, inRas.extent.YMin)
        cellSize = inRas.meanCellWidth

        # 使用我们确定的output_nodata_val来创建栅格
        newRaster = arcpy.NumPyArrayToRaster(final_arr, lowerLeft, cellSize, value_to_nodata=output_nodata_val)

        newRaster.save(outRecalssRaster_full_path)
        arcpy.AddMessage("成功保存输出栅格: {}".format(outRecalssRaster_full_path))
        return outRecalssRaster_full_path

    except Exception as e:
        arcpy.AddError("处理过程中发生错误:")
        arcpy.AddError(str(e))
        return None

def main():
    parser = argparse.ArgumentParser(description='对栅格图像作自然断点法分类脚本')
    parser.add_argument('--in_RasterFullPath', required=True, help='输入的栅格影像路径')
    parser.add_argument('--nb_class', type=int, required=True, help='危险性等级数量')
    parser.add_argument('--nb_times', type=int, required=True, help='自然断点法采样次数')
    parser.add_argument('--sample_size', type=int, required=True, help='自然断点法每次采样的样本数量')
    parser.add_argument('--out_ReclassRasterFullPath', required=True, help='输出的栅格影像路径')
    parser.add_argument('--nodata_val', type=int, required=True, help='输出的栅格影像中NoData的值')
    args = parser.parse_args()
    nb_class = int(args.nb_class)
    nb_times = int(args.nb_times)
    sample_size = int(args.sample_size)
    nodata_val = int(args.nodata_val)
    jenkspy_Natural_Breaks_sample(args.in_RasterFullPath, args.out_ReclassRasterFullPath,
                                  nb_class, nb_times=nb_times, sample_size=sample_size,
                                  nodata_val=nodata_val)

def tt():
    in_RasterFullPath = r"C:\Users\DDTwo\Documents\ArcGIS\Projects\gwlsaAngDa\RasterUnits20220905_csv_202511140849.tif"
    out_ReclassRasterFullPath = r"C:\Users\DDTwo\Documents\ArcGIS\Projects\gwlsaAngDa\gwlsaAngDa.gdb\RasterUnits20220905_csv_202511140849_202511140944NaturalBreaks"
    nb_class = 4
    nb_times = 50
    sample_size = 5000
    nodata_val = 9999
    jenkspy_Natural_Breaks_sample(in_RasterFullPath, out_ReclassRasterFullPath,
                                  nb_class, nb_times=nb_times, sample_size=sample_size,
                                  nodata_val=nodata_val)

if __name__ == '__main__':
    main()
    # tt()


